﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;
using Ast.Core.Events;
using Ast.Core.Services.Extensions;
using Ast.Core.WinApi.DbgHelp;
using Ast.Framework;
using Castle.Core.Logging;

namespace Ast.Core.Services
{
    public interface ISymbolService : ISingletonDependency
    {
        void Initialize();

        ulong LoadModule(string imageName, string moduleName, ulong baseOfDll);
        bool UnloadModule(ulong baseOfDll);
        bool UnloadAllModules();
        bool ListModules(Action<string, ulong> procModuleFunc);

        DbgHelpStructure.SymbolInfo? LoadSymbol(ulong address);
        DbgHelpStructure.SymbolInfo? LoadSymbol(string name);
        bool ListSymbols(ulong address, string mask, Func<DbgHelpStructure.SymbolInfo, bool> procSymbolFunc);
        bool SearchSymbols(ulong baseOfDll, DbgHelpConstant.SymTagEnum symTag, string mask, Func<DbgHelpStructure.SymbolInfo, bool> procSymbolFunc);

        DbgHelpStructure.SymbolInfo? GetSymbolType(ulong baseOfDll, string name);
        T GetSymbolTypeInfo<T>(ulong baseOfDll, uint typeId, DbgHelpConstant.ImagehlpSymbolTypeInfo getType);
        DbgHelpStructure.SymbolInfo? GetSymbolTypeInfo(ulong baseOfDll, uint typeId);
        uint[] GetSymbolTypeInfoChildren(ulong baseOfDll, uint typeId);
        string GetSymbolTypeInfoName(ulong baseOfDll, uint typeId);

        ulong FindAddressSpace(uint size);
    }

    public class SymbolService : ISymbolService
    {
        private readonly IDbgHelpService _dbgHelpService;
        private readonly ISymbolEvents _symbolEvents;

        private DbgHelpNative.SymRegisterCallbackProc64 _symbolNotificationCallback; // keep the reference

        public SymbolService(IDbgHelpService dbgHelpService, ISymbolEvents symbolEvents)
        {
            _dbgHelpService = dbgHelpService;
            _symbolEvents = symbolEvents;

            Logger = NullLogger.Instance;
        }

        public ILogger Logger { get; set; }

        public void Initialize()
        {
            _dbgHelpService.AcquireLock();

            // options
            var options = _dbgHelpService.SymGetOptions() | DbgHelpConstant.SymOption.SYMOPT_DEBUG | DbgHelpConstant.SymOption.SYMOPT_CASE_INSENSITIVE;
            _dbgHelpService.SymSetOptions(options);

            // notification
            _symbolNotificationCallback = SymbolNotificationCallback;
            _dbgHelpService.SymRegisterCallback64(_symbolNotificationCallback);

            _dbgHelpService.ReleaseLock();

            Logger.Info("initialized symbol service");
        }

        private bool SymbolNotificationCallback(IntPtr hProcess, DbgHelpConstant.SymActionCode actionCode, ulong callbackData, ulong userContext)
        {
            var dataPtr = new IntPtr((int) callbackData);
            if (actionCode == DbgHelpConstant.SymActionCode.CBA_DEFERRED_SYMBOL_LOAD_START)
            {
                var info = Marshal.PtrToStructure<DbgHelpStructure.ImagehlpDeferredSymbolLoad64>(dataPtr);
                Logger.InfoFormat("symbol loading has started for {0}", info.FileName);
            }
            else if (actionCode == DbgHelpConstant.SymActionCode.CBA_DEFERRED_SYMBOL_LOAD_PARTIAL)
            {
                var info = Marshal.PtrToStructure<DbgHelpStructure.ImagehlpDeferredSymbolLoad64>(dataPtr);
                Logger.InfoFormat("symbol load has partially completed for {0}", info.FileName);
            }
            else if (actionCode == DbgHelpConstant.SymActionCode.CBA_DEFERRED_SYMBOL_LOAD_COMPLETE)
            {
                var info = Marshal.PtrToStructure<DbgHelpStructure.ImagehlpDeferredSymbolLoad64>(dataPtr);
                Logger.InfoFormat("symbol load has completed for {0}", info.FileName);
            }
            else if (actionCode == DbgHelpConstant.SymActionCode.CBA_DEFERRED_SYMBOL_LOAD_FAILURE)
            {
                var info = Marshal.PtrToStructure<DbgHelpStructure.ImagehlpDeferredSymbolLoad64>(dataPtr);
                Logger.InfoFormat("symbol load has failed for {0}", info.FileName);
            }
            else if (actionCode == DbgHelpConstant.SymActionCode.CBA_DEBUG_INFO)
            {
                var info = Marshal.PtrToStringAnsi(dataPtr) ?? "";
                if (!info.StartsWith("SYMSRV:  Waiting for the server to respond to a request.") &&
                    !info.StartsWith("SYMSRV:  Successfully received a response from the server."))
                {
                    Logger.Info(info);
                }
            }
            return false;
        }

        public ulong LoadModule(string imageName, string moduleName, ulong baseOfDll)
        {
            var isLoaded = false;
            var retVal = 0ul;

            _dbgHelpService.AcquireLock();

            // checking loaded symbols
            moduleName = moduleName ?? Path.GetFileNameWithoutExtension(imageName);
            _dbgHelpService.SymEnumerateModules64((name, dll, context) =>
            {
                if (name == moduleName)
                {
                    retVal = dll;
                    return false;
                }
                return true;
            });

            if (retVal == 0)
            {
                // loading
                retVal = _dbgHelpService.SymLoadModuleEx(imageName, moduleName, baseOfDll);
                if (retVal != 0) isLoaded = true;
            }

            _dbgHelpService.ReleaseLock();

            // events
            if (isLoaded)
            {
                _symbolEvents.OnLoadedModule(imageName, moduleName, retVal);
            }

            return retVal;
        }

        public bool UnloadModule(ulong baseOfDll)
        {
            _dbgHelpService.AcquireLock();
            var isOk = _dbgHelpService.SymUnloadModule64(baseOfDll);
            _dbgHelpService.ReleaseLock();

            // events
            if (isOk)
            {
                _symbolEvents.OnUnloadedModule(baseOfDll);
            }

            return isOk;
        }

        public bool UnloadAllModules()
        {
            var isOk = true;
            var moduleBases = new List<ulong>();

            _dbgHelpService.AcquireLock();
            _dbgHelpService.SymEnumerateModules64((name, dll, context) =>
            {
                moduleBases.Add(dll);
                return true;
            });
            foreach (var moduleBase in moduleBases)
            {
                isOk = UnloadModule(moduleBase); // unload one by one
                if (!isOk) break;
            }
            _dbgHelpService.ReleaseLock();

            return isOk;
        }

        public bool ListModules(Action<string, ulong> procModuleFunc)
        {
            _dbgHelpService.AcquireLock();
            var isOk = _dbgHelpService.SymEnumerateModules64((name, dll, context) =>
            {
                procModuleFunc(name, dll);
                return true;
            });
            _dbgHelpService.ReleaseLock();

            return isOk;
        }

        public DbgHelpStructure.SymbolInfo? LoadSymbol(ulong address)
        {
            _dbgHelpService.AcquireLock();
            var symbolInfo = DbgHelpStructure.SymbolInfo.Create();
            var isFound = _dbgHelpService.SymFromAddr(address, ref symbolInfo);
            _dbgHelpService.ReleaseLock();
            return isFound ? symbolInfo : null as DbgHelpStructure.SymbolInfo?;
        }

        public DbgHelpStructure.SymbolInfo? LoadSymbol(string name)
        {
            _dbgHelpService.AcquireLock();
            var symbolInfo = DbgHelpStructure.SymbolInfo.Create();
            var isFound = _dbgHelpService.SymFromName(name, ref symbolInfo);
            _dbgHelpService.ReleaseLock();
            return isFound ? symbolInfo : null as DbgHelpStructure.SymbolInfo?;
        }

        public bool ListSymbols(ulong address, string mask, Func<DbgHelpStructure.SymbolInfo, bool> procSymbolFunc)
        {
            _dbgHelpService.AcquireLock();
            var isOk = _dbgHelpService.SymEnumSymbols(address, mask, (ref DbgHelpStructure.SymbolInfo info, uint size, IntPtr context) => procSymbolFunc(info));
            _dbgHelpService.ReleaseLock();
            return isOk;
        }

        public bool SearchSymbols(ulong baseOfDll, DbgHelpConstant.SymTagEnum symTag, string mask, Func<DbgHelpStructure.SymbolInfo, bool> procSymbolFunc)
        {
            _dbgHelpService.AcquireLock();
            var isOk = _dbgHelpService.SymSearch(baseOfDll, 0, symTag, mask, 0, (ref DbgHelpStructure.SymbolInfo info, uint size, IntPtr context) => procSymbolFunc(info));
            _dbgHelpService.ReleaseLock();
            return isOk;
        }

        public DbgHelpStructure.SymbolInfo? GetSymbolType(ulong baseOfDll, string name)
        {
            _dbgHelpService.AcquireLock();
            var symbolInfo = DbgHelpStructure.SymbolInfo.Create();
            var isFound = _dbgHelpService.SymGetTypeFromName(baseOfDll, name, ref symbolInfo);
            _dbgHelpService.ReleaseLock();
            return isFound ? symbolInfo : null as DbgHelpStructure.SymbolInfo?;
        }

        public T GetSymbolTypeInfo<T>(ulong baseOfDll, uint typeId, DbgHelpConstant.ImagehlpSymbolTypeInfo getType)
        {
            var retVal = default(T);
            _dbgHelpService.AcquireLock();
            using (var memoryBlock = LocalMemoryBlock.Allocate<T>())
            {
                if (_dbgHelpService.SymGetTypeInfo(baseOfDll, typeId, getType, memoryBlock.Address))
                {
                    retVal = memoryBlock.ReadAs<T>();
                }
            }
            _dbgHelpService.ReleaseLock();
            return retVal;
        }

        public DbgHelpStructure.SymbolInfo? GetSymbolTypeInfo(ulong baseOfDll, uint typeId)
        {
            _dbgHelpService.AcquireLock();
            var symbolInfo = DbgHelpStructure.SymbolInfo.Create();
            var isFound = _dbgHelpService.SymFromIndex(baseOfDll, typeId, ref symbolInfo);
            _dbgHelpService.ReleaseLock();
            return isFound ? symbolInfo : null as DbgHelpStructure.SymbolInfo?;
        }

        public uint[] GetSymbolTypeInfoChildren(ulong baseOfDll, uint typeId)
        {
            var childrenCount = GetSymbolTypeInfo<int>(baseOfDll, typeId, DbgHelpConstant.ImagehlpSymbolTypeInfo.TI_GET_CHILDRENCOUNT);
            if (childrenCount == 0) return null;

            using (var children = LocalMemoryBlock.Allocate(Marshal.SizeOf<TiFindchildrenParams>() + Marshal.SizeOf<UInt32>()*childrenCount))
            {
                var childrenParams = new TiFindchildrenParams {Count = (uint) childrenCount};
                Marshal.StructureToPtr(childrenParams, children.Address, true);

                _dbgHelpService.AcquireLock();
                var isOk = _dbgHelpService.SymGetTypeInfo(baseOfDll, typeId, DbgHelpConstant.ImagehlpSymbolTypeInfo.TI_FINDCHILDREN, children.Address);
                _dbgHelpService.ReleaseLock();

                return isOk
                    ? children.ReadAsBytes(Marshal.SizeOf<UInt32>()*childrenCount, Marshal.SizeOf<TiFindchildrenParams>()).ReadArrayAs<uint>(childrenCount)
                    : null;
            }
        }

        public string GetSymbolTypeInfoName(ulong baseOfDll, uint typeId)
        {
            var namePtr = GetSymbolTypeInfo<IntPtr>(baseOfDll, typeId, DbgHelpConstant.ImagehlpSymbolTypeInfo.TI_GET_SYMNAME);
            if (namePtr == IntPtr.Zero) return null;

            var name = Marshal.PtrToStringUni(namePtr);
            Marshal.FreeHGlobal(namePtr); // free local
            return name;
        }

        public ulong FindAddressSpace(uint size)
        {
            // first space
            var start = 0x2400000ul; // starting address
            size = ((size - 1)/0x1000)*0x1000 + 0x1000; // making sure it is times of 0x1000
            var end = start + size;

            // loading existing modules
            var usedAddresses = new List<AddressSpace>();
            _dbgHelpService.AcquireLock();
            _dbgHelpService.SymEnumerateModules64((name, dll, context) =>
            {
                var module = DbgHelpStructure.ImagehlpModule64.Create();
                _dbgHelpService.SymGetModuleInfo64(dll, ref module);
                if (dll >= start)
                {
                    usedAddresses.Add(new AddressSpace(dll, module.ImageSize));
                }
                return true;
            });
            _dbgHelpService.ReleaseLock();

            // last space
            usedAddresses.Add(new AddressSpace(ulong.MaxValue, 0));

            // work out the first available address space
            foreach (var addressSpace in usedAddresses.OrderBy(x => x.Start))
            {
                if (end < addressSpace.Start)
                {
                    return start;
                }
                start = ((addressSpace.End - 1)/0x1000)*0x1000 + 0x1000 + 0x1000; // additional 0x1000 gap
                end = start + size;
            }

            // unavailable
            return 0;
        }

        public class AddressSpace
        {
            public AddressSpace(ulong start, uint size)
            {
                Start = start;
                Size = size;
            }

            public ulong Start { get; set; }

            public ulong End
            {
                get { return Start + Size; }
            }

            public uint Size { get; set; }
        }

        [StructLayout(LayoutKind.Sequential, CharSet = CharSet.Auto, Pack = 1)]
        public struct TiFindchildrenParams
        {
            public uint Count;
            public uint Start;
            // public uint ChildId;
        }
    }
}